# -*- coding: utf-8 -*-
# %%
'''
from Crypto.Protocol.SecretSharing import Shamir
from binascii import hexlify
from binascii import unhexlify
'''
from Crypto.PublicKey import DSA
from CryptoAlgorithm.SecretSharing.sharing import PlaintextToHexSecretSharer
import os
from CryptoAlgorithm.ElgamalSig import elgamal
from CryptoAlgorithm import ElgamalSig
import time
from FileOperate import DataProcess


# %%

class ElGamal_Shamir(object):

    def shamir_encrypt(self, in_file, num, sum, bin_file):
        """
        secret sharing encrypting with shamir
        :param in_file: key file needed encrypt
        :param num: you can descrypt the file with "num" files
        :param sum: the sum of shamir files
        :param bin_file: the result of keyfile encrypted
        """

        def share_fn(x):
            """
            创建shamir密钥文件名
            """
            x = bin_file + x + "_shamir_key.bin"
            return x

        sh_filen = [str(x) for x in range(1, sum + 1)]
        fi_filen = list(map(share_fn, sh_filen)) # ['1_shamir_key.bin', '2_shamir_key.bin', '5_shamir_key.bin']
        for i in fi_filen:
            if os.path.exists(i) is True:
                os.remove(i)
        with open(in_file, 'rb') as kf:
            kf_file = kf.readline()
            while (kf_file):
                share_text = PlaintextToHexSecretSharer. \
                    split_secret(bytes.decode(kf_file), num, sum)
                i = 0
                for text in share_text:
                    with open(fi_filen[i], 'a+') as fi_file:
                        fi_file.write(text + "\n")
                    i = i + 1
                kf_file = kf.readline()

        #os.remove(key_file)



    def shamir_descrypt(self, out_file, sum, bin_file):
        """
        secret sharing descrypting with shamir
        :param out_file: the complete DSA key file name that you want
        :param sum: the share file sum number
        :param bin_file: the result of keyfile encrypted
        """

        def share_fn(x):
            """
            创建shamir密钥文件名
            """
            x = bin_file + x + "_shamir_key.bin"
            return x

        sh_filen = [str(x) for x in range(1, sum + 1)]
        fi_filen = list(map(share_fn, sh_filen))

        if os.path.exists(out_file) is True:
            os.remove(out_file)

        try:
            sum_temp = []
            for sh_file in fi_filen:
                if os.path.exists(sh_file):
                    with open(sh_file, 'r') as f:
                        temp = []
                        for line in f.readlines():
                            temp.append(line.rstrip("\n"))
                        sum_temp.append(temp)
            for i in range(len(sum_temp[0])):
                text = []
                for j in range(len(sum_temp)):
                    text.append(sum_temp[j][i])
                text = PlaintextToHexSecretSharer.recover_secret(text)
                text = str.encode(text)
                with open(out_file, 'ab') as f:
                    f.write(text)

        except:
            raise IOError
        '''
        for fn in fi_filen:
            if os.path.exists(fn) is True:
                os.remove(fn)
        '''


if __name__ == "__main__":
    start = time.time()
    #message = "Hello"
    model_path = '../data/ML/sourceModel/second/user0/model0.h5'
    model_data_list = DataProcess().readModel(model_path)
    message_str = ElgamalSig.dataListToJson(model_data_list)  # <class 'str'>

    user_path = "../data/socket/signature/user0/"
    if not os.path.exists(user_path):
        os.mkdir(user_path)

    pri_path = user_path + "private_key.bin"
    pub_path = user_path + "public_key.pem"
    pri_path_temp = user_path + "pri_key.bin"

    key = elgamal().produceKey(user_path)

    signature = elgamal().sign(message_str, key, user_path)

    ElGamal_Shamir().shamir_encrypt(pri_path, 3, 5, user_path)
    ElGamal_Shamir().shamir_descrypt(pri_path_temp, 3, user_path)


    # 验签
    with open(user_path + 'sign.bin', 'rb') as f:
        sign_new = bytearray(f.read())
    with open(pri_path_temp,'rb')as x:
        private_key = DSA.import_key(x.read())
    new_pub_key_temp = private_key.publickey().export_key()  # b'jshdfjhajfhkaj-===='
    new_pub_key = DSA.import_key(new_pub_key_temp)  # <DsaKey @0x1c99cb4f6a0 y,g,p(2048),q>

    elgamal().verify(message_str, sign_new, new_pub_key)
    end = time.time()
    # 统计训练时间
    total = end - start
    showtime = time.strftime('%M:%S', time.localtime(total))
    print('密钥分割与合并运行时间:' + showtime)